/***************************************************************************
 *   Copyright (C) 2009-2010 by Mayank Lahiri                              *
 *   mlahiri@gmail.com                                                     *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/
#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string>
#include <string.h>
#include <time.h>
#include <map>
#include <limits.h>
#include <vector>
#include <unistd.h>
#include <sys/types.h>
#include <algorithm>
#include "netcollab.h"
#include "readers.h"
#include "mt19937ar.h"
#include "tokenizer.h"
#include "utils.h"

#define CURRENT_VERSION 1.5

using namespace std;

//
// jump table for various file format readers
//
void *jump_table[]  =   {   (void*) "enron",    (void*)reader_enron,
                            (void*) "dirpair",  (void*)reader_dirpair,
                            (void*) "undirpair",(void*)reader_undirpair,
                            NULL };

//
// globals
//
vector<INTERACTION> data;
char in_file[256];
char out_file[256];
char in_format[32];
bool directed                           = false;        // are interaction directed in this dataset?
unsigned nv                             = 0;            // number of individuals
unsigned ne                             = 0;            // number of interactions
unsigned random_seed                    = 0;            // random seed for PRNG
unsigned trial                          = 0;            // current trial
unsigned method                         = 0;            // perturbation method
unsigned time_pert                      = 0;            // time perturbation window in seconds
unsigned TRIALS                         = 50;
int NEW_HDF                             = 0;            // generate a new HDF for every trial?
int HDF_LEN                             = -1;           // bit length of HDF function (-1 == randomly chosen in [256,1024])
int POPDIST                             = 0;
int VERSION                             = 0;
unsigned RECORD_INTERVAL                = 1000;         // intervals (in # of iterations) at which to record population stats
RESULT_BY_ITER *iter_result             = NULL;
RESULT_BY_NODE *node_result             = NULL;


// command line parameters
CMDLINE_PARAMETERS cmd[]	={
	{ "-i", CMDLINE_STRING, in_file, 1, 255, 1, "input file" },
	{ "-o", CMDLINE_STRING, out_file, 1, 255, 1, "output file stem"},
	{ "-t", CMDLINE_STRING, in_format, 1, 32, 1, "input file format ('enron','dirpair','undirpair')"},
	{ "-m", CMDLINE_INTEGER, &method, 0, 4, 1, "method (0=normal, 1=random topology, 2=full shuffle, 3=time-perturbed)" },
	{ "-seed", CMDLINE_INTEGER, &random_seed, 0, INT_MAX, 0, "random seed (otherwise chosen randomly)"},
	{ "-pert", CMDLINE_INTEGER, &time_pert,  0, 999999999, 0, "time perturbation window for method = 3 (in seconds)" },
	{ "-trials", CMDLINE_INTEGER, &TRIALS, 0, 9999999, 0, "trials"},
	{ "-newhdf", CMDLINE_BINARY, &NEW_HDF, 0, 1, 0, "generate new HDF function for each trial"},
	{ "-hdflen", CMDLINE_INTEGER, &HDF_LEN, 1,9999999, 0, "HDF function length (otherwise chosen randomly in [256,1024]" },
	{ "-popdist", CMDLINE_BINARY, &POPDIST, 0, 1, 0, "dump population fitness distributions at increments of (1/5th) of total iterations"},
	{ "-sample", CMDLINE_INTEGER, &RECORD_INTERVAL, 1, 9999999, 0, "interactions between records of population statistics"},
	{ "-h", CMDLINE_HELP, NULL, 0, 0, 0, "this help message"},
	{ "-v", CMDLINE_BINARY, &VERSION, 0, 1, 0, "print version"},
	{ "", 0, NULL, 0, 0, 0, "" }
};

//
// forward declarations
//
void fullShuffle();
void partialShuffle();
void writeResults();
int  version();
void dumpPopDist(int, GA_hdf &);

//
// program entry point
//
int main(int argc, char *argv[])	{
    //
    // setup
    //
    if(argc > 1 && argv[1][0] == '-' && argv[1][1] == 'v')
        return version();
	random_seed = sqrt(time(NULL)) - clock() + sqrt(getpid());
	FILE *fp;
	if((fp = fopen("/dev/urandom","rb")))
	    if(fread(&random_seed, sizeof(unsigned), 1, fp))
            fclose(fp);
	if(!ParseCommandLine(argc, argv, cmd))
		return -1;
    if(method == 3 && time_pert == 0)   {
        fprintf(stderr, "Please specify a value for -pert if using method 3.\n");
        return -2;
    }
    init_genrand(random_seed);
    printf("# NetCollab v%.1f copyright (c) 2009-2010 by Mayank Lahiri -- mlahiri@gmail.com. Released under the terms of the GNU GPL\n", CURRENT_VERSION);
    printf("# Initial random seed is %u\n", random_seed);

    // try to read the input file
    void **s = jump_table;
    while(*s)  {
        if(strcmp((char*)*s, in_format)==0) {
            s++;
            // call the file reader function in jump_table
            ne = ((int (*)(char*,vector<INTERACTION>&,bool&,unsigned &))*s) (in_file, data, directed, nv);
            if(ne < 1)
                return -3;
            break;
        }
        s += 2;
    }
    if(!*s) {
        fprintf(stderr, "Unrecognized format '%s' for input file '%s'\n", in_format, in_file);
        return -4;
    }
    printf("# Read %u %sdirected interactions between %u individuals from '%s'\n", ne, directed?"":"un", nv, in_file);
    fflush(stdout);
    assert(ne == data.size());
    RECORD_INTERVAL = data.size() / 100;

    // allocate space to hold the results
    iter_result = new RESULT_BY_ITER[102];
    node_result = new RESULT_BY_NODE[nv];

    // create the GA (use an arbitrary length genome with the default range)
	GA_hdf ga(nv, HDF_LEN);

	// start the trials
	vector<double> itr_mean, itr_max;
	for(trial = 1; trial <= TRIALS; trial++)	{
        // reset the GA
		if(NEW_HDF)
			ga.hdf_generate(HDF_LEN);
        ga.randomize();

		// prepare for this trial
		printf("# Starting trial %d with seed %u, initial max fit %g\n", trial, random_seed+trial-1, ga.pop_max()); fflush(stdout);
		init_genrand(random_seed+trial-1);
		itr_max.clear();
		itr_mean.clear();

        // perturb the interactions if asked for
        if(method == 2)
			// full shuffle
			fullShuffle();
        else
            if(method == 3)
				// constrained shuffle
				partialShuffle();

        // *********************************************************************
        // BEGINNING OF TRIAL
		// *********************************************************************
        // iterate through all the interactions (FINALLY!)
        int pdistmarker = 0;
		for(unsigned r = 0; r < data.size(); r++)	{
		    unsigned v1 = data[r].from;
		    unsigned v2 = data[r].to;

		    // forget about the actual vertices if we want random topology
		    if(method == 1) {
		        v1 = genrand_real2() * nv;
		        do {
		            v2 = genrand_real2() * nv;
		        } while(v1 == v2);
		    }

		    // do the interaction
		    ga.interact(v1, v2, directed);

		    // update our summary stats every RECORD_INTERVAL interactions
		    if(r%RECORD_INTERVAL == 0) {
				itr_mean.push_back(ga.pop_mean());
				itr_max.push_back(ga.pop_max());
		    }

            // do we want population fitness dumps?
			if(POPDIST && (r%(data.size()/5)==0))	{
				printf("+--> %u%% interactions processed, dumping fitness distribution\n", (unsigned)(r*100/data.size()));
				dumpPopDist(pdistmarker++, ga);
			}

            // keep the console user happy, will ya?
			if((r+1) % 1000000 == 0){
				printf("# Processed %d million interactions\n", 1+r/1000000);
				fflush(stdout);
			}
		}
		// *********************************************************************
		// END OF TRIAL
		// *********************************************************************

        // normalize and update (end-of-trial) statistics
        assert(itr_max.size() == itr_mean.size());
        assert(itr_mean.size() <= 102);
        double mad = 0, rmsd = 0;
        for(unsigned i = 0; i < itr_mean.size(); i++)   {
            itr_mean[i]         /= ga.pop_max();
            itr_max[i]          /= ga.pop_max();
            mad                 += fabs(iter_result[i].mean/(trial==1?1:trial-1) - (iter_result[i].mean+itr_mean[i])/trial);
            rmsd                += pow(iter_result[i].mean/(trial==1?1:trial-1) - (iter_result[i].mean+itr_mean[i])/trial, 2.0);

            iter_result[i].mean     += itr_mean[i];
            iter_result[i].mean_sq  += itr_mean[i] * itr_mean[i];
            iter_result[i].max      += itr_max[i];
            iter_result[i].max_sq   += itr_max[i] * itr_max[i];
        }
        mad  /= itr_mean.size();
        rmsd  = sqrt(rmsd/itr_mean.size());

		// record final (end-of-trial) node statistics
		double finalfit_mad = 0, finalfit_rmsd = 0, finalfit_max = 0;
		for(unsigned n = 0; n < nv; n++)	{
			GA_NODE *nod = ga.node(n);
			assert(nod->obj / ga.pop_max() <= 1.0);
			double diff                 =  fabs(node_result[n].final_obj/(trial==1?1:trial-1) - (nod->obj/ga.pop_max()+node_result[n].final_obj)/trial);
			finalfit_mad                += diff;
			finalfit_rmsd               += pow(diff, 2.0);
			finalfit_max                =  diff>finalfit_max ? diff : finalfit_max;

			node_result[n].final_obj    += nod->obj / ga.pop_max();
			node_result[n].final_obj_sq += pow(nod->obj/ga.pop_max(), 2.0);
			node_result[n].improve      += nod->improvements;
			node_result[n].improve_sq   += pow(nod->improvements, 2.0);
			node_result[n].contacts     += nod->contacts;
			node_result[n].contacts_sq  += nod->contacts * nod->contacts;
			node_result[n].neighbors    += nod->neighbors.size();
			node_result[n].neighbors_sq += nod->neighbors.size() * nod->neighbors.size();
		}
		finalfit_mad /= nv;
		finalfit_rmsd/= nv;

		// write results so far
		writeResults();
		printf("TRIAL %d mean_fit_norm %.4f max_fitness %.1f MEAN_DELTA_FITCURVE[ mad %.3f rmsd %.3f ] MEAN_DELTA_FINALFIT[ mad %.3f rmsd %.3f max %.3f ]\n",
                trial, ga.pop_mean()/ga.pop_max(), ga.pop_max(), mad, rmsd, finalfit_mad, finalfit_rmsd, finalfit_max);
	}

	delete[] node_result;
	delete[] iter_result;
	return EXIT_SUCCESS;
}

//
// Fisher-Yates shuffle of the interaction list
//
void fullShuffle()	{
	printf("# Running full shuffle\n");
	for(int i = (int)data.size(); i > 1; )	{
		unsigned k = genrand_real2() * i;
		i--;
		INTERACTION t = data[i-1];
		data[i-1] = data[k];
		data[k] = t;
	}
}

//
// perturbs each interaction in time by an amount chosen
// uniformly from a fixed window
//
void partialShuffle()	{
	printf("# Running partial shuffle\n");
    for(unsigned i = 0; i < data.size(); i++)
        data[i].perturb = genrand_real1()*2.0*time_pert - time_pert;
    sort(data.begin(), data.end());
}


void writeResults()	{
	char fbuf[300];
	sprintf(fbuf, "%s.iter", out_file);
	FILE *fpOut = fopen(fbuf, "w+");
	if(!fpOut)	{
		fprintf(stderr, "WARNING: Cannot open %s for writing results\n", fbuf);
		return;
	}
	fprintf(fpOut, "# In_file '%s' Out_file '%s' method '%d' seed '%u' time_perturb '%d' trials '%d' record_interval '%d' newhdf '%d' hdflen '%d'\n",
					in_file, out_file, method, random_seed, time_pert, trial, RECORD_INTERVAL, NEW_HDF, HDF_LEN);
	fprintf(fpOut, "iter pop_mean dev_pop_mean pop_max dev_pop_max\n");
	for(unsigned r = 0; r < data.size()/RECORD_INTERVAL; r++)
		fprintf(fpOut, "%d %f %f %f %f\n", r*RECORD_INTERVAL,
						(double)iter_result[r].mean/trial,sqrt((double)iter_result[r].mean_sq/trial - pow((double)iter_result[r].mean/trial,2.0)),
						(double)iter_result[r].max/trial,sqrt((double)iter_result[r].max_sq/trial - pow((double)iter_result[r].max/trial,2.0)));
	fclose(fpOut);

	sprintf(fbuf, "%s.node", out_file);
	fpOut = fopen(fbuf, "w+");
	if(!fpOut)	{
		fprintf(stderr, "WARNING: Cannot open %s for writing results\n", fbuf);
		return;
	}
	fprintf(fpOut, "# In_file '%s' Out_file '%s' method '%d' seed '%u' time_perturb '%d' trials '%d' record_interval '%d' newhdf '%d' hdflen '%d'\n",
					in_file, out_file, method, random_seed, time_pert, trial, RECORD_INTERVAL, NEW_HDF, HDF_LEN);
	fprintf(fpOut, "ind_number ind_label AFF AFFdev contacts contactsdev improve improvedev neighbors neighborsdev\n");
	for(unsigned r = 0; r < nv; r++)    {
		fprintf(fpOut, "%d %s %.4f %.4f %.1f %.2f %.1f %.2f %.1f %.2f\n", r, reader_retranslate(r).c_str(),
                        (double)(node_result[r].final_obj/trial), sqrt(node_result[r].final_obj_sq/trial - pow((double)node_result[r].final_obj/trial,2.0)),
						(double)(node_result[r].contacts/trial) , sqrt(node_result[r].contacts_sq/trial  - pow((double)node_result[r].contacts/trial,2.0) ),
						(double)(node_result[r].improve/trial)  , sqrt(node_result[r].improve_sq/trial   - pow((double)node_result[r].improve/trial,2.0)),
						(double)(node_result[r].neighbors/trial), sqrt(node_result[r].neighbors_sq/trial - pow((double)node_result[r].neighbors/trial,2.0))
						 );
	}
    fflush(fpOut);
	fclose(fpOut);
}

int version()  {
    printf("Netcollab v%.1f, compiled %s %s\nCopyright 2010 Mayank Lahiri, released under the GNU GPL.\n", CURRENT_VERSION, __DATE__, __TIME__);
    return 0;
}

vector<double> inc_pop_fit_dist[6];
vector<double> temp_fit_dist[5];

void dumpPopDist(int marker, GA_hdf &ga)   {
    if(temp_fit_dist[0].size() != nv)   {
        // initialize
        for(int w = 0; w < 6; w++)  {
            if(w!=5)    {
                temp_fit_dist[w].clear();
                for(unsigned i = 0; i < nv; i++)
                    temp_fit_dist[w].push_back(0.0);
            }
            for(unsigned i = 0; i < nv; i++)
                inc_pop_fit_dist[w].push_back(0.0);
        }
    }
    assert(temp_fit_dist[0].size() == nv);

    if(marker < 5)
        for(unsigned i = 0; i < nv; i++)
            temp_fit_dist[marker][i] = ga.node(i)->obj;
    else    {
        // final marker, normalize earlier distributions by pop max
        // and write to file
        for(int m = 0; m < 5; m++)
            for(unsigned i = 0; i < nv; i++)    {
                assert(temp_fit_dist[m][i] / ga.pop_max() <= 1.0);
                inc_pop_fit_dist[m][i] += (temp_fit_dist[m][i] / ga.pop_max());
            }
        for(unsigned i = 0; i < nv; i++)
            inc_pop_fit_dist[5][i] += ga.node(i)->obj / ga.pop_max();

        // now write it to disk
        for(int m = 0; m < 6; m++)  {
            char buf[256];
            sprintf(buf, "%s.popdist.%d", out_file, m);
            FILE *fp = fopen(buf, "w+");
            if(!fp)	{
                fprintf(stderr, "WARNING: Cannot open %s for writing results\n", buf);
                return;
            }
            for(unsigned i = 0; i < nv; i++)
                fprintf(fp, "%.3f\n", inc_pop_fit_dist[m][i]/trial);
            fclose(fp);
        }
    }
}
